# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
rStar-Math 主程序入口

这是 rStar-Math 数学推理系统的主要入口文件，负责：
1. 解析命令行参数和配置文件
2. 初始化推理求解器 (Solver)
3. 选择搜索算法 (MCTS 或束搜索)
4. 批量处理数学问题并保存结果

主要处理流程：
- 加载问题答案数据集 (JSON/JSONL格式)
- 根据配置创建相应的搜索代理 (MCTS/BS)
- 使用 Solver 对问题进行推理求解
- 将结果保存为 JSONL 文件

使用示例：
python main.py --qaf eval_data/GSM8K_test.json --custom_cfg config/sft_eval_mcts.yaml --model_dir your-model-path
"""

from __future__ import annotations  # 启用延迟注解评估，支持前向引用类型注解
import os  # 操作系统接口，用于文件路径操作和环境变量设置
import json  # JSON数据格式处理，用于加载和保存数据文件
import torch  # PyTorch深度学习框架，用于模型推理和张量计算
import argparse  # 命令行参数解析器，处理程序启动参数
from tqdm import tqdm  # 进度条显示库，用于显示批处理进度
from datetime import datetime  # 日期时间处理，用于生成时间戳文件名
from omegaconf import OmegaConf  # 配置管理库，用于加载和合并YAML配置文件
from rstar_deepthink.agents import BS, MCTS, MCTSWoCode, MCTSCaption  # 搜索算法代理：束搜索(BS)、蒙特卡洛树搜索(MCTS)、无代码MCTS(MCTSWoCode)和Caption扩展MCTS(MCTSCaption)
from rstar_deepthink.solver import Solver  # 核心求解器，负责协调LLM和搜索算法进行推理
from rstar_deepthink.config import BaseConfig  # 基础配置类，定义系统默认配置参数

# 设置 PyTorch 线程数，控制并行计算资源
torch.set_num_threads(12)
# 禁用 tokenizers 并行化，避免与多进程冲突
os.environ["TOKENIZERS_PARALLELISM"] = "false"

def load_qaf(filename: str, config=None):
    """
    加载问题答案文件 (Question-Answer File)
    
    支持两种数据格式：
    1. JSON格式：包含问题和答案的结构化数据
    2. JSONL格式：每行一个JSON对象的文件
    
    参数:
        filename: 数据文件路径，支持 .json 和 .jsonl 扩展名
        config: 配置对象，用于多模态数据处理和数据集配置（可选）
        
    返回:
        list: 包含 {"question": str, "answer": str} 格式的数据列表
        
    异常:
        ValueError: 当文件格式不被支持时抛出
    """
    if filename.endswith(".json"):
        # 处理标准JSON格式文件
        with open(filename, "r") as f:
            data = json.load(f)
        # 某些数据集将实际数据嵌套在 "example" 键下
        if "example" in data:
            data = data["example"]
    elif filename.endswith(".jsonl"):
        # 处理JSONL格式文件（每行一个JSON对象）
        data = []
        with open(filename, "r") as f:
            lines = f.readlines()
        for line in lines:
            data.append(json.loads(line))
    else:
        raise ValueError(f"Unrecognized file format: {filename}")
    
    # 应用数据集处理配置
    if config:
        # 随机打乱数据集（如果启用）
        if config.dataset_shuffle:
            import random
            if config.dataset_seed is not None:
                random.seed(config.dataset_seed)
            random.shuffle(data)
            print(f"数据集已随机打乱，种子: {config.dataset_seed}")
        
        # 应用起始索引和最大样本数限制
        original_len = len(data)
        start_idx = max(0, config.dataset_start_idx)
        
        if config.dataset_max_samples > 0:
            end_idx = min(len(data), start_idx + config.dataset_max_samples)
        else:
            end_idx = len(data)
        
        data = data[start_idx:end_idx]
        print(f"数据集处理范围: [{start_idx}:{end_idx}]，原始数据 {original_len} 条，处理 {len(data)} 条")
    
    # 处理多模态数据
    if config and config.multimodal_enabled:
        data = _process_multimodal_data(data, config)
    
    return data

def _process_multimodal_data(data, config):
    """
    处理多模态数据，标准化字段映射和图像路径
    
    参数:
        data: 原始数据列表
        config: 配置对象
        
    返回:
        list: 处理后的数据列表
    """
    processed_data = []
    
    for item in data:
        processed_item = {}
        
        # 标准化字段映射
        processed_item["question"] = item.get(config.question_field, item.get("question", ""))
        processed_item["answer"] = item.get(config.answer_field, item.get("answer", ""))
        
        # 处理图像字段
        if config.image_field in item:
            image_path = item[config.image_field]
            # 如果是相对路径，转换为绝对路径
            if not os.path.isabs(image_path):
                image_path = os.path.join(config.image_root, image_path)
            processed_item[config.image_field] = image_path
        
        # 保留其他字段
        for key, value in item.items():
            if key not in [config.question_field, config.answer_field, config.image_field]:
                processed_item[key] = value
        
        processed_data.append(processed_item)
    
    return processed_data

def batch(iterable, n=-1):
    """
    将可迭代对象分批处理的生成器函数
    
    用于控制内存使用和并行处理效率，将大数据集分成小批次处理。
    
    参数:
        iterable: 要分批的可迭代对象（如数据列表）
        n: 批次大小，-1或小于等于0时表示处理全部数据
        
    生成:
        每次生成一个批次的数据子集
        
    示例:
        data = [1, 2, 3, 4, 5]
        for batch_data in batch(data, 2):
            print(batch_data)  # [1, 2], [3, 4], [5]
    """
    l = len(iterable)
    if n <= 0:
        n = l  # 如果批次大小无效，则处理全部数据
    for ndx in range(0, l, n):
        yield iterable[ndx: min(ndx + n, l)]

def parse_args():
    """
    解析命令行参数
    
    配置 rStar-Math 推理系统的各项参数，包括模型路径、数据文件、配置文件等。
    
    返回:
        argparse.Namespace: 解析后的命令行参数对象
        
    主要参数说明:
        --custom_cfg: 自定义配置文件路径，覆盖默认配置
        --qaf: 问题答案数据文件路径
        --model_dir: 策略模型（LLM）目录路径
        --reward_model_dir: PPM奖励模型目录路径
        --save_in_model: 自定义输出文件名（可选）
        --output_dir: 输出目录路径，默认为 output
    """
    args = argparse.ArgumentParser()
    # 配置文件路径，默认使用MCTS评估配置
    args.add_argument('--custom_cfg', type=str, default="config/sft_eval_mcts.yaml",
                     help="自定义配置文件路径，用于覆盖默认配置参数")
    # 问题答案数据文件
    args.add_argument("--qaf", type=str, default="", 
                     help="问题答案文件路径（JSON或JSONL格式）")
    # 策略模型路径
    args.add_argument('--model_dir', type=str, default="",
                     help="策略模型（LLM）目录路径或HuggingFace模型名称") 
    # 奖励模型路径
    args.add_argument('--reward_model_dir', type=str, default="",
                     help="PPM过程偏好模型目录路径") 
    # 自定义输出文件名
    args.add_argument('--save_in_model', type=str, default="",
                     help="自定义输出文件名，不指定则自动生成时间戳文件名")
    # 输出目录路径
    args.add_argument('--output_dir', type=str, default="output",
                     help="输出目录路径，默认为 output")
    args = args.parse_args()
    return args


if __name__ == '__main__':
    # === 1. 参数解析和配置初始化 ===
    args = parse_args()

    # 创建基础配置对象
    config = OmegaConf.structured(BaseConfig)
    # 如果指定了自定义配置文件，则合并配置
    if args.custom_cfg:
        custom_config = OmegaConf.load(args.custom_cfg)
        config = OmegaConf.merge(config, custom_config)
    # 解析配置中的占位符和引用
    config = OmegaConf.create(OmegaConf.to_yaml(config, resolve=True))
    
    # 命令行参数覆盖配置文件中的相应设置
    if args.model_dir:
        config.model_dir = args.model_dir
    if args.reward_model_dir:
        config.reward_model_dir = args.reward_model_dir
    print(config)

    # === 2. 数据加载和模型初始化 ===
    # 提取模型版本信息用于文件命名
    llm_version = os.path.basename(config.model_dir.rstrip("/"))

    # 加载问题答案数据集（支持多模态）
    data = load_qaf(args.qaf, config)
    # 创建推理求解器，初始化LLM和奖励模型
    solver = Solver(config=config)

    # === 3. 搜索算法选择 ===
    # 根据配置选择相应的搜索算法
    if config.mode == "mcts":
        agent = MCTS  # 蒙特卡洛树搜索
    elif config.mode == "mcts_wo_code":
        agent = MCTSWoCode  # 无代码模式的蒙特卡洛树搜索
    elif config.mode == "mcts_caption":
        agent = MCTSCaption  # Caption 扩展的蒙特卡洛树搜索
    elif config.mode == "bs":
        agent = BS    # 束搜索
    else:
        raise NotImplementedError(f"不支持的搜索模式: {config.mode}")
    
    # 如果使用奖励模型，将其版本信息添加到文件名中
    if args.reward_model_dir:
        llm_version += "." + args.reward_model_dir.split("/")[-1]
        
    # === 4. 输出文件配置 ===
    # 确保输出目录存在
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 创建本次运行的专用目录
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    run_name = f"{config.mode}_{llm_version}_{timestamp}"
    run_dir = os.path.join(args.output_dir, run_name)
    os.makedirs(run_dir, exist_ok=True)
    
    # 创建子目录
    results_dir = os.path.join(run_dir, "results")
    rollout_dir = os.path.join(run_dir, "rollouts") 
    logs_dir = os.path.join(run_dir, "logs")
    os.makedirs(results_dir, exist_ok=True)
    os.makedirs(rollout_dir, exist_ok=True)
    os.makedirs(logs_dir, exist_ok=True)
    
    # 生成输出文件名：数据文件.搜索模式.模型版本.时间戳.jsonl
    base_filename = f"{os.path.basename(args.qaf)}.{config.mode}.{llm_version}.{timestamp}.jsonl"
    saved_jsonl_file = os.path.join(results_dir, base_filename)
    
    # 创建运行信息文件
    run_info = {
        "timestamp": timestamp,
        "config_file": args.custom_cfg,
        "data_file": args.qaf,
        "model_dir": config.model_dir,
        "mode": config.mode,
        "data_count": len(data),
        "config": OmegaConf.to_container(config, resolve=True)
    }
    with open(os.path.join(run_dir, "run_info.json"), "w", encoding='utf-8') as f:
        json.dump(run_info, f, ensure_ascii=False, indent=2)
    
    print(f"\n📁 本次运行目录: {run_dir}")
    print(f"📊 结果输出: {saved_jsonl_file}")
    print(f"🔄 Rollout记录: {rollout_dir}")
    print(f"📝 日志记录: {logs_dir}")
    
    # 如果指定了自定义输出文件名，则使用自定义名称
    if args.save_in_model:
        if os.path.isabs(args.save_in_model):
            # 绝对路径，直接使用
            saved_jsonl_file = args.save_in_model + '.jsonl'
        else:
            # 相对路径，放在结果目录下
            saved_jsonl_file = os.path.join(results_dir, args.save_in_model + '.jsonl')
        saved_jsonl_file_dir = os.path.dirname(saved_jsonl_file)
        os.makedirs(saved_jsonl_file_dir, exist_ok=True)
        
    # === 5. 批量推理处理 ===
    # 以追加模式打开输出文件，支持断点续传
    with open(saved_jsonl_file, "a+", encoding='utf-8') as writer:
        # 按配置的批次大小处理数据
        for cur_data in tqdm(batch(data, config.batch_size), desc="Main Processing"):
            # 为当前批次的每个问题创建搜索代理
            agents = [agent(config=config, question=d["question"], ground_truth=str(d["answer"])) 
                      for d in cur_data]
            
            # 设置多模态数据到根节点
            if config.multimodal_enabled:
                for i, a in enumerate(agents):
                    if config.image_field in cur_data[i]:
                        a.root.multimodal_data = {
                            config.image_field: cur_data[i][config.image_field]
                        }
            
            # 使用求解器进行推理，返回每个问题的推理结果
            jsonlines = solver.solve(agents, saved_jsonl_file, cur_data)
            # 将推理结果写入输出文件
            for d in cur_data:
                question = d["question"]
                d["rstar"] = jsonlines[question]  # 添加rStar推理结果
                writer.write(json.dumps(d, ensure_ascii=False) + '\n')
                writer.flush()  # 确保数据及时写入磁盘
